#!/usr/bin/env python

# emailfilter.py
# Kevin Altis
# http://altis.pycs.net/
# altis@semi-retired.com
# created:  2003-08-22
# modified: 2003-08-24
# version: 0.2
#
# originally created to auto-delete SoBig.F virus attachments from POP3 accounts
# the script should run on any platform with Python 2.2.2 or later
# http://www.python.org/

# you may also want to look at an alternative method posted by Alex
# http://www.tranzoa.com/extras.htm#remove_email_virus

# module documentation is at:
# http://docs.python.org/lib/module-poplib.html
# http://docs.python.org/lib/module-email.html

import sys
# we need Python 2.2.2 or later to get version 2 of the email package
assert sys.version_info >= (2, 2, 2)

import os
import getpass
import poplib
import email
from email import Header
import ConfigParser

USAGE = """usage:
    python emailfilter.py
    python emailfilter.py filename
    python emailfilter.py server username password

The last option is best if you don't want to store your username and 
password in a plaintext .ini file. If you don't want to use the 
command-line options, then the .ini file should look like:

[Account]
server=your_server
username=your_username
password=your_password"""

DEFAULTINI = 'account.ini'
DEBUGLEVEL = 0


def getAttachments(msg):
    attachments = []
    # try MIME
    counter = 0
    for part in msg.walk():
        # multipart/* are just containers
        partType = part.get_type()
        if part.get_main_type() == 'multipart':
            continue
        filename = part.get_filename()
        ### this part is new
        ### contributed by Andrew Dalke
        ### for MS Update virus
        if filename is None:
            if partType in ["application/x-msdownload",
                            "audio/x-wav", "audio/x-midi"]:
                # Kill with extreme prejudice
                filename = "virus.exe"
                          
        #### end of new part
        if filename is not None:
            if not filename:
                ext = mimetypes.guess_extension(part.get_type())
                if not ext:
                    # Use a generic bag-of-bits extension
                    ext = '.bin'
                filename = 'part-%03d%s' % (counter, ext)
            #print counter, partType, filename
            attachments.append({'part':counter, 'type':partType, 'filename':filename})
        counter += 1
    return attachments

# delete "failure notice" messages too?
# these are returned messages where the virus
# payload is no longer an attachment
# alternatively could search for strings
# in the payload and delete if found
# a typical attachment looks like:
"""
Content-Type: application/octet-stream;
    name="document_all.pif"
Content-Transfer-Encoding: base64
Content-Disposition: attachment;
    filename="document_all.pif"
"""
def isVirus(message):
    # SoBig.F has headers like the ones
    # below prefixed by Re: 
    # but this handler doesn't currently match based on Subject
    subjectLines = [
        'My details', 
        'Your details', 
        'Your application', 
        'Wicked screensaver', 
        'That movie', 
        'Approved', 
        'Details', 
        'Thank you!',
        ]

    # these are typical extensions
    # for spreading viruses
    extensions = [
        'bat',
        'com',
        'exe',
        'pif',
        'reg',
        'scr',
        'vbs',
        'wsh',
        ]

    # could flag based on a variety of criteria
    # including Subject line, but for now
    # we'll only return True if the message
    # has a payload that ends in pif, exe, etc.

    virus = False
    attachments = getAttachments(message)
    if len(attachments) > 0:
        for a in attachments:
            print "    Attachment:", a['filename']
            base, ext = os.path.splitext(a['filename'])
            if ext[1:].lower() in extensions:
                virus = True
    return virus

# expects an open mailbox
def deleteViruses(mailbox):
    numMessages, totalSize = mailbox.stat()
    messages = []
    for i in range(1, numMessages + 1):
        #message = mailbox.retr(i+1)
        # use top to avoid the message being marked as read
        messageTuple = mailbox.top(i, -1)
        message = email.message_from_string("\n".join(messageTuple[1]))
        # if for some reason you need to collect the messages
        # uncomment the line below
        # messages.append(message)
        encodedSubject = message.get('subject')
        # handle internationalized subject line
        # http://docs.python.org/lib/module-email.Header.html
        subject, encoding = Header.decode_header(encodedSubject)[0]
        #print i, "Subject:", subject, encoding
        if encoding is None or encoding == 'iso-8859-1':
            print i, "Subject:", subject
        else:
            print i, "Subject:", encodedSubject
        if isVirus(message):
            print "    VIRUS"
            mailbox.dele(i)

def processMailbox(server, username, password, debuglevel=DEBUGLEVEL):
    mailbox = poplib.POP3(server)
    try:
        mailbox.user(username)
        mailbox.pass_(password)
        mailbox.set_debuglevel(debuglevel)
        numMessages, totalSize = mailbox.stat()
        print "Messages:   ", numMessages
        print "Total Bytes:", totalSize
        print "=" * 78

        deleteViruses(mailbox)
    finally:
        mailbox.quit()


if __name__ == '__main__':
    if len(sys.argv) == 4:
        server = sys.argv[1]
        username = sys.argv[2]
        password = sys.argv[3]
    else:
        if len(sys.argv) == 2:
            path = sys.argv[1]
        else:
            path = DEFAULTINI
        parser = ConfigParser.ConfigParser()
        parser.read(path)
        server = parser.get('Account', 'server')
        username = parser.get('Account', 'username')
        password = parser.get('Account', 'password')
    processMailbox(server, username, password)
